// ----------------------------------------------------------------------------
//  ServerlessLLM
//  Copyright (c) ServerlessLLM Team 2024
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//
//  You may obtain a copy of the License at
//
//              http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
//  ----------------------------------------------------------------------------
#pragma once

#include <string>
#include <vector>

// CANN headers - these will be available when building with CANN toolkit
#ifdef USE_CANN
#include "acl/acl.h"
#endif

class CannMemoryPool {
 public:
  CannMemoryPool(int device_count, size_t size_per_device);
  CannMemoryPool(const CannMemoryPool&) = delete;
  CannMemoryPool& operator=(const CannMemoryPool&) = delete;
  ~CannMemoryPool();

  int Allocate(size_t size, int device_id, void*& ptr,
               std::string& handle);  // CANN uses string handles
  int Deallocate(int device_id, void* ptr);

 private:
  int device_count_;
  size_t size_per_device_;
  std::vector<void*> pool_;
  std::vector<std::string> handles_;  // CANN IPC handles as strings
  std::vector<size_t> offsets_;
};